# TODO: 导入必要的库和模块
import numpy as np
import pickle
import gradio as gr
from PIL import Image

# TODO: 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
try:
    with open('best_knn_model.pkl', 'rb') as f:
        model = pickle.load(f)
    print("模型加载成功")
except FileNotFoundError:
    print("错误：请先运行optimal_knn.py生成模型文件")
    exit(1)

# TODO: 定义预测函数，这个函数将用于Gradio接口进行预测
def preprocess_image(image):
    """将用户绘制的图像转换为模型需要的格式"""
    if image is None:
        return None
    
    # 转换为灰度图并调整大小为8x8（与训练数据一致）
    img = Image.fromarray(image).convert('L')
    img = img.resize((8, 8), Image.LANCZOS)
    img_array = np.array(img)
    
    # 反转颜色（训练数据是白色数字+黑色背景）
    img_array = 255 - img_array
    
    # 归一化到0-16范围（与digits数据集一致）
    img_array = (img_array / 255) * 16
    
    # 展平为1D数组
    return img_array.flatten().reshape(1, -1)

def predict_digit(image):
    """预测用户绘制的数字"""
    if image is None:
        return "请先绘制一个数字"
    
    # 预处理图像
    processed_img = preprocess_image(image)
    if processed_img is None:
        return "图像处理失败"
    
    # 模型预测
    prediction = model.predict(processed_img)
    return f"预测结果：{prediction[0]}"

# TODO: 创建Gradio接口，这个接口将用于用户输入和显示预测结果
with gr.Blocks(title="手写数字识别") as interface:
    gr.Markdown("# KNN手写数字识别")
    gr.Markdown("在下方画板绘制0-9的数字，点击提交获取预测结果")
    
    with gr.Row():
        sketchpad = gr.Sketchpad(shape=(200, 200), label="绘制区域")
        output_text = gr.Textbox(label="预测结果")
    
    with gr.Row():
        submit_btn = gr.Button("提交")
        clear_btn = gr.Button("清除")
    
    # 绑定事件
    submit_btn.click(
        fn=predict_digit,
        inputs=sketchpad,
        outputs=output_text
    )
    clear_btn.click(
        fn=lambda: None,
        inputs=None,
        outputs=sketchpad,
        queue=False
    )

# TODO: 启动Gradio接口，用户可以通过这个接口进行交互
if __name__ == "__main__":
    interface.launch()
